# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Force filed"""
import os
from typing import Union, List
import numpy as np
from numpy import ndarray
import torch
from torch import Tensor
from torch.nn import ModuleList

from .energy import EnergyCell, get_energy_cell
from .potential import PotentialCell
from ..data.forcefield import get_forcefield
from ..system import Molecule
from ..function import get_arguments, Units, Length, get_tensor


THIS_PATH = os.path.abspath(__file__)
BUILTIN_FF_PATH = THIS_PATH.replace('potential/forcefield.py', 'data/forcefield/')


class ForceFieldBase(PotentialCell):
    r"""Base class for the potential energy of classical force filed. It is a subclass of `PotentialCell`.

        A `ForceFieldBase` object contains multiple `EnergyCell` objects. The last dimension of its output Tensor
        is equal to the number of `EnergyCell` objects it contains.

    Args:

        Energy (Union[EnergyCell, List[EnergyCell]]):
                                List of `EnergyCell` objects. Default: None

        cutoff (Union[float, Length, Tensor]):  Cutoff distance. Default: None

        exclude_index (Union[Tensor, ndarray, List[int]]):
                                Array of indexes of atoms that should be excluded from neighbour list.
                                The shape of the tensor is `(B, A, Ex)`. The data type is int.
                                Default: None

        length_unit (str):      Length unit. If None is given, it will be assigned with the global length unit.
                                Default: None

        energy_unit (str):      Energy unit. If None is given, it will be assigned with the global energy unit.
                                Default: None

        use_pbc (bool):         Whether to use periodic boundary condition.

    Returns:

        energy (Tensor):    Tensor of shape `(B, E)`. Data type is float.

    Supported Platforms:

        ``CPU`` ``GPU``

    Symbols:

        B:  Batchsize, i.e. number of walkers in simulation.
        E:  Number of energy terms.

    """

    def __init__(self,
                 energy: Union[EnergyCell, List[EnergyCell]] = None,
                 cutoff: Union[float, Length, Tensor] = None,
                 exclude_index: Union[Tensor, ndarray, List[int]] = None,
                 length_unit: str = None,
                 energy_unit: str = None,
                 use_pbc: bool = None,
                 name: str = 'potential',
                 **kwargs,
                 ):

        super().__init__(
            length_unit=length_unit,
            energy_unit=energy_unit,
            use_pbc=use_pbc,
            name=name,
        )
        self._kwargs = get_arguments(locals(), kwargs)

        if isinstance(cutoff, Length):
            cutoff = cutoff(self.units)

        if cutoff is not None:
            self.cutoff = get_tensor(cutoff, dtype=torch.float32)

        self._exclude_index = self._check_exclude_index(exclude_index)

        self._num_energies = 0
        self._energy_index = {}
        self.energies: List[EnergyCell] = None
        self.output_unit_scale = 1
        self.set_energies(energy)

        self.energy_scale = 1

    def set_energy_scale(self, scale: Tensor):
        """set energy scale"""
        scale = get_tensor(scale, dtype=torch.float32)
        if scale.ndim not in (0, 1):
            raise ValueError('The rank of energy scale must be 0 or 1.')
        if scale.shape[-1] != self._num_energies and scale.shape[-1] != 1:
            raise ValueError(f'The dimension of energy scale must be equal to '
                             f'the dimension of energy {self._num_energies} or 1, '
                             f'but got: {scale.shape[-1]}')
        self.energy_scale = scale
        return self

    def set_energies(self, energy: EnergyCell):
        """set energy"""
        if energy is None:
            self.output_unit_scale = 1
            return self
        if isinstance(energy, EnergyCell):
            self._num_energies = 1
            self.energies = ModuleList([energy])
        elif isinstance(energy, list):
            self._num_energies = len(energy)
            self.energies = ModuleList(energy)
        elif isinstance(energy, ModuleList):
            self._num_energies = len(energy)
            self.energies = energy
        else:
            raise TypeError(f'The type of energy must be EnergyCell or list but got: {type(energy)}')

        self._energy_names = []
        self._energy_index = {}
        index = 0
        if self.energies is not None:
            for i in range(self._num_energies):
                name = self.energies[i].name
                self._energy_index[name] = index
                self._energy_names.append(name)
                index += 1

        self.set_unit_scale()
        return self

    def append_energy(self, energy: EnergyCell):
        """append energy terms"""
        if not isinstance(energy, EnergyCell):
            raise TypeError(f'The type of energy must be EnergyCell or list but got: {type(energy)}')

        self.energies.append(energy)
        self._energy_names.append(energy.name)
        self._energy_index[energy.name] = self._num_energies
        self._num_energies += 1
        self.set_unit_scale()

        return self

    def set_unit_scale(self) -> Tensor:
        """set unit scale"""
        if self.energies is None:
            self.output_unit_scale = 1
            return self
        output_unit_scale = []
        for i in range(self._num_energies):
            self.energies[i].set_input_unit(self.units)
            scale = self.energies[i].convert_energy_to(self.units)
            output_unit_scale.append(scale)
        self.output_unit_scale = get_tensor(output_unit_scale, dtype=torch.float32)
        return self

    def set_units(self, length_unit: str = None, energy_unit: str = None, units: Units = None):
        """set units"""
        super().set_units(length_unit=length_unit, energy_unit=energy_unit, units=units)
        self.set_unit_scale()
        return self

    def set_pbc(self, use_pbc: bool = None):
        """set whether to use periodic boundary condition."""
        for i in range(self._num_energies):
            self.energies[i].set_pbc(use_pbc)
        return self

    def set_cutoff(self, cutoff: float = None, unit: str = None):
        """set cutoff distance"""
        if cutoff is None:
            energy_cutoff = []
            for i in range(self._num_energies):
                cutoff = self.energies[i].cutoff
                if cutoff is not None:
                    cutoff = self.units.length(cutoff, self.energies[i].length_unit)
                    energy_cutoff.append(cutoff)
            if energy_cutoff:
                cutoff = max(energy_cutoff)
                super().set_cutoff(cutoff)
        else:
            super().set_cutoff(cutoff, unit)
            for i in range(self._num_energies):
                self.energies[i].set_cutoff(self.cutoff, self.length_unit)
        return self

    def forward(self,
                  coordinate: Tensor,
                  neighbour_index: Tensor = None,
                  neighbour_mask: Tensor = None,
                  neighbour_vector: Tensor = None,
                  neighbour_distance: Tensor = None,
                  pbc_box: Tensor = None
                  ):
        r"""Calculate potential energy.

        Args:
            coordinate (Tensor):            Tensor of shape (B, A, D). Data type is float.
                                            Position coordinate of atoms in system.
            neighbour_index (Tensor):       Tensor of shape (B, A, N). Data type is int.
                                            Index of neighbour atoms. Default: None
            neighbour_mask (Tensor):        Tensor of shape (B, A, N). Data type is bool.
                                            Mask for neighbour atoms. Default: None
            neighbour_vector (Tensor):       Tensor of shape (B, A, N, D). Data type is bool.
                                            Vectors from central atom to neighbouring atoms.
            neighbour_distance (Tensor):    Tensor of shape (B, A, N). Data type is float.
                                            Distance between neighbours atoms. Default: None
            pbc_box (Tensor):               Tensor of shape (B, D). Data type is float.
                                            Tensor of PBC box. Default: None

        Returns:
            potential (Tensor): Tensor of shape (B, E). Data type is float.

        Symbols:
            B:  Batchsize, i.e. number of walkers in simulation.
            A:  Number of atoms.
            N:  Maximum number of neighbour atoms.
            D:  Spatial dimension of the simulation system. Usually is 3.
            E:  Number of energy terms.

        """
        #pylint: disable=unused-argument

        if self.energies is None:
            return torch.zeros((coordinate.shape[0], 1), dtype=coordinate.dtype, device=coordinate.device)

        energies = []
        for i in range(self._num_energies):
            energy = self.energies[i](
                coordinate=coordinate,
                neighbour_index=neighbour_index,
                neighbour_mask=neighbour_mask,
                neighbour_vector=neighbour_vector,
                neighbour_distance=neighbour_distance,
                pbc_box=pbc_box
            )
            energies.append(energy.to(torch.float32))

        # (B,E)
        energies = torch.cat(energies, dim=-1) * self.output_unit_scale * self.energy_scale
        return energies

    def print_info(self) -> dict:
        """Return basic information about the force field
        
        Returns:
            dict: Dictionary containing force field information
        """
        info = {
            "Number of energy terms": self._num_energies,
            "Energy terms": self._energy_names if self._num_energies > 0 else None,
            "Cutoff": self.cutoff.item() if self.cutoff is not None else 'No cutoff',
            "Energy unit": self.units.energy_unit,
            "Length unit": self.units.length_unit
        }
        return info


class ForceField(ForceFieldBase):
    r"""Base class for the potential energy of classical force filed. It is a subclass of `ForceFieldBase`.

        A `ForceField` object contains multiple `EnergyCell` objects. The last dimension of its output Tensor
        is equal to the number of `EnergyCell` objects it contains.

    Args:

        system (Molecule):      Molecule object.

        parameters (Union[dict, str, List[Union[dict, str]]]):
                                Parameters of force field. It can be a dictionary, a string of the path of
                                parameter file, or a list of them.

        cutoff (float):         Cutoff distance. Default: None

        rebuild_system (bool):  Whether to rebuild the system. Default: True

        length_unit (str):      Length unit. If None is given, it will be assigned with the global length unit.
                                Default: None

        energy_unit (str):      Energy unit. If None is given, it will be assigned with the global energy unit.
                                Default: None

        name (str):             Name of force field. Default: "potential"

    Returns:

        energy (Tensor):    Tensor of shape `(B, E)`. Data type is float.

    Supported Platforms:

        ``CPU`` ``GPU``

    Symbols:

        B:  Batchsize, i.e. number of walkers in simulation.
        E:  Number of energy terms.

    """

    def __init__(self,
                 system: Molecule,
                 parameters: Union[dict, str, List[Union[dict, str]]],
                 cutoff: float = None,
                 rebuild_system: bool = True,
                 length_unit: str = None,
                 energy_unit: str = None,
                 name: str = 'potential',
                 **kwargs,
                 ):

        super().__init__(
            cutoff=cutoff,
            exclude_index=self.get_exclude_index(system),
            length_unit=length_unit,
            energy_unit=energy_unit,
            name=name,
            **kwargs,
        )

        self._kwargs = get_arguments(locals(), kwargs)
        
        # 保存力场名称
        self.ff_name = parameters if isinstance(parameters, str) else None

        parameters, template = get_forcefield(parameters)
        if rebuild_system:
            if template is None:
                print('[WARNING] No template in parameters! Cannot rebuild the system!')
            else:
                for res in system.residue:
                    tplt = template.get(res.name)
                    if tplt is not None:
                        res.build_atom_type(tplt)
                        res.build_atom_charge(tplt)
                system.build_system()

        energy = []
        for energy_name, energy_params in parameters.items():
            energy_ = get_energy_cell(cls_name=energy_name,
                                        system=system,
                                        parameters=energy_params,
                                        exclude_index=self._exclude_index,
                                        **kwargs
                                        )
            if energy_ is not None:
                energy.append(energy_)

        self.set_energies(energy)
        self.set_cutoff(self.cutoff, self.length_unit)

    @staticmethod
    def get_exclude_index(system: Molecule) -> Tensor:
        """get exclude index"""
        if system.bonds is None:
            return None

        num_atoms = system.num_atoms
        bonds = system.bonds[0].detach().cpu().numpy()
        angles = None if system.angles is None else system.angles.detach().cpu().numpy()
        dihedrals = None if system.dihedrals is None else system.dihedrals.detach().cpu().numpy()
        improper = None if system.improper_dihedrals is None else system.improper_dihedrals.detach().cpu().numpy()

        excludes_ = []
        for i in range(num_atoms):
            bond_excludes = bonds[np.where(
                np.isin(bonds, i).sum(axis=1))[0]].flatten()
            this_excludes = bond_excludes

            if angles is not None:
                angle_excludes = angles[
                    np.where(np.isin(angles, i).sum(axis=1))[0]
                ].flatten()
                this_excludes = np.append(this_excludes, angle_excludes)

            if dihedrals is not None:
                dihedral_excludes = dihedrals[
                    np.where(np.isin(dihedrals, i).sum(axis=1))[0]
                ].flatten()
                this_excludes = np.append(this_excludes, dihedral_excludes)
            if improper is not None:
                idihedral_excludes = improper[
                    np.where(np.isin(improper, i).sum(axis=1))[0]
                ].flatten()
                this_excludes = np.append(this_excludes, idihedral_excludes)

            this_excludes = np.unique(this_excludes)
            excludes_.append(this_excludes[np.where(
                this_excludes != i)[0]].tolist())
        padding_length = 0
        for i in range(num_atoms):
            padding_length = max(padding_length, len(excludes_[i]))
        excludes = np.empty((num_atoms, padding_length))
        for i in range(num_atoms):
            excludes[i] = np.pad(
                np.array(excludes_[i]),
                (0, padding_length - len(excludes_[i])),
                mode="constant",
                constant_values=num_atoms,
            )
        return get_tensor(excludes[None, :], dtype=torch.int32)

    def print_info(self) -> dict:
        """Return basic information about the force field
        
        Returns:
            dict: Dictionary containing force field information
        """
        info = super().print_info()
        if self.ff_name is not None:
            info["Force field"] = self.ff_name
        return info
